package vip.zhenzicheng.websocket.websocket.config.interceptor;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import javax.servlet.http.HttpServletRequest;
import java.util.Map;

/**
 * 在握手的过程中对token进行判断(拦截器)，也可以通过继承{@link AbstractWebSocketHandler}及其子处理器来对socket连接事件处理
 */
@Slf4j
public class TokenInterceptor implements HandshakeInterceptor {

    /**
     * WebSocket自定义传输协议
     */
    String WEBSOCKET_CUSTOM_PROTOCOL = "Sec-WebSocket-Protocol";

    /**
     * 握手阶段对token进行解析，websocket通过自定义协议"sec-websocket-protocol"模拟header传参
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        log.info("即将握手");
        if (request instanceof ServletServerHttpRequest) {
            HttpServletRequest httpServletRequest = ((ServletServerHttpRequest) request).getServletRequest();
            String token = httpServletRequest.getHeader(WEBSOCKET_CUSTOM_PROTOCOL);
            log.info("\n\n解析到自定义协议为:{}\n\n", token);
            // Assert.hasLength(token, "token不存在!");
            // request.getHeaders().put("token", Collections.singletonList(token));
            // 将token保存在webSocketSession中,方便在handler中直接获取,放在header中可能在某些地方获取不到
            // attributes.put("token", token);
            // 重新保存自定义协议到响应头，防止报错
            ((ServletServerHttpResponse) response).getServletResponse().addHeader(WEBSOCKET_CUSTOM_PROTOCOL, token);
        }
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        log.info("握手完成");
    }
}
